from functools import wraps
from contextlib import contextmanager
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.orm import class_mapper
from hetong.config import config
engine = create_engine(
        config.SQLALCHEMY_DATABASE_URI,
        max_overflow=10,  # 超过连接池大小外最多创建的连接
        pool_size=40,  # 连接池大小
        pool_timeout=30,  # 池中没有线程最多等待的时间，否则报错
        pool_recycle=-1,  # 多久之后对线程池中的线程进行一次连接的回收（重置）
        echo=True,
        pool_pre_ping=True  # 每次使用连接前检查连接是否可用

    )
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
@contextmanager
def session_scope() -> Session:
    """上下文管理器用于自动获取 Session, 避免错误"""
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except:
        session.rollback()
        raise
    finally:
        session.close()


def with_session(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        with session_scope() as session:
            try:
                result = f(session, *args, **kwargs)
                session.commit()
                return result
            except:
                session.rollback()
                raise

    return wrapper


def get_db() -> SessionLocal:
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def get_db0() -> SessionLocal:
    db = SessionLocal()
    return db

def asdict(obj):
    return {col.key: getattr(obj, col.key) for col in class_mapper(obj.__class__).columns}

